# -*- coding: utf-8 -*-
"""
阿里云百炼Embedding模型兼容OpenAI接口规范
支持将原有OpenAI应用迁移至阿里云百炼，只需调整三个参数：
- base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
- api_key: 阿里云百炼API Key
- model: text-embedding-v4 或 text-embedding-v3
"""

import os
import sys
from typing import List, Optional, Union
from openai import OpenAI
import numpy as np

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(project_root)

from config_reader import get_bailian_config


class DashScopeEmbedding:
    """
    阿里云百炼Embedding模型兼容OpenAI接口规范
    支持text-embedding-v4和text-embedding-v3模型
    """
    
    def __init__(self, 
                 api_key: Optional[str] = None,
                 base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1",
                 model: str = "text-embedding-v4",
                 dimensions: Optional[int] = None,
                 encoding_format: str = "float"):
        """
        初始化阿里云百炼Embedding客户端
        
        Args:
            api_key: 阿里云百炼API Key，如果为None则从配置文件读取
            base_url: API基础URL，默认为华北2（北京）地域
            model: 模型名称，支持text-embedding-v4和text-embedding-v3
            dimensions: 向量维度，仅text-embedding-v3和text-embedding-v4支持
            encoding_format: 编码格式，默认为float
        """
        # 从配置文件获取API Key（如果未提供）
        if api_key is None:
            bailian_config = get_bailian_config()
            api_key = bailian_config['token']
        
        if not api_key:
            raise ValueError("请提供阿里云百炼API Key或在config.ini中设置BAILIAN_API_TOKEN")
        
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.dimensions = dimensions
        self.encoding_format = encoding_format
        
        # 初始化OpenAI客户端
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url
        )
    
    def get_embedding(self, text: str) -> List[float]:
        """
        获取单个文本的嵌入向量
        
        Args:
            text: 输入文本
            
        Returns:
            嵌入向量列表
        """
        try:
            # 构建请求参数
            params = {
                "model": self.model,
                "input": text,
                "encoding_format": self.encoding_format
            }
            
            # 如果指定了维度，添加到参数中
            if self.dimensions is not None:
                params["dimensions"] = self.dimensions
            
            response = self.client.embeddings.create(**params)
            return response.data[0].embedding
            
        except Exception as e:
            print(f"获取嵌入向量时发生错误: {e}")
            raise
    
    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        获取多个文本的嵌入向量
        
        Args:
            texts: 输入文本列表
            
        Returns:
            嵌入向量列表的列表
        """
        try:
            # 构建请求参数
            params = {
                "model": self.model,
                "input": texts,
                "encoding_format": self.encoding_format
            }
            
            # 如果指定了维度，添加到参数中
            if self.dimensions is not None:
                params["dimensions"] = self.dimensions
            
            response = self.client.embeddings.create(**params)
            return [item.embedding for item in response.data]
            
        except Exception as e:
            print(f"获取嵌入向量时发生错误: {e}")
            raise
    
    def cosine_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
        """
        计算两个嵌入向量的余弦相似度
        
        Args:
            embedding1: 第一个嵌入向量
            embedding2: 第二个嵌入向量
            
        Returns:
            余弦相似度值（-1到1之间）
        """
        try:
            # 转换为numpy数组
            vec1 = np.array(embedding1)
            vec2 = np.array(embedding2)
            
            # 计算余弦相似度
            dot_product = np.dot(vec1, vec2)
            norm1 = np.linalg.norm(vec1)
            norm2 = np.linalg.norm(vec2)
            
            if norm1 == 0 or norm2 == 0:
                return 0.0
            
            similarity = dot_product / (norm1 * norm2)
            return float(similarity)
            
        except Exception as e:
            print(f"计算余弦相似度时发生错误: {e}")
            return 0.0
    
    def similarity_search(self, query_text: str, candidate_texts: List[str], top_k: int = 5) -> List[tuple]:
        """
        在候选文本中搜索与查询文本最相似的文本
        
        Args:
            query_text: 查询文本
            candidate_texts: 候选文本列表
            top_k: 返回最相似的前k个结果
            
        Returns:
            包含(文本, 相似度分数)的元组列表，按相似度降序排列
        """
        try:
            # 获取查询文本的嵌入向量
            query_embedding = self.get_embedding(query_text)
            
            # 获取所有候选文本的嵌入向量
            candidate_embeddings = self.get_embeddings(candidate_texts)
            
            # 计算相似度
            similarities = []
            for i, candidate_embedding in enumerate(candidate_embeddings):
                similarity = self.cosine_similarity(query_embedding, candidate_embedding)
                similarities.append((candidate_texts[i], similarity))
            
            # 按相似度降序排序
            similarities.sort(key=lambda x: x[1], reverse=True)
            
            # 返回前top_k个结果
            return similarities[:top_k]
            
        except Exception as e:
            print(f"相似度搜索时发生错误: {e}")
            return []


def main():
    """
    示例用法
    """
    try:
        # 创建DashScope Embedding实例
        embedding = DashScopeEmbedding(
            model="text-embedding-v4",
            dimensions=1024  # 指定向量维度
        )
        
        # 示例文本
        texts = [
            "狗是伟大的伴侣，以其忠诚和友好而闻名。",
            "猫是独立的宠物，通常喜欢自己的空间。",
            "金鱼是初学者的流行宠物，需要相对简单的护理。",
            "鹦鹉是聪明的鸟类，能够模仿人类的语言。",
            "兔子是社交动物，需要足够的空间跳跃。"
        ]
        
        print("=== 阿里云百炼Embedding模型测试 ===")
        print(f"使用模型: {embedding.model}")
        print(f"向量维度: {embedding.dimensions}")
        print()
        
        # 测试单个文本嵌入
        print("1. 单个文本嵌入测试:")
        query_text = "咖啡猫"
        query_embedding = embedding.get_embedding(query_text)
        print(f"查询文本: {query_text}")
        print(f"嵌入向量维度: {len(query_embedding)}")
        print(f"嵌入向量前5个值: {query_embedding[:5]}")
        print()
        
        # 测试相似度搜索
        print("2. 相似度搜索测试:")
        results = embedding.similarity_search(query_text, texts, top_k=3)
        print(f"查询: {query_text}")
        print("最相似的文本:")
        for i, (text, similarity) in enumerate(results, 1):
            print(f"  {i}. 相似度: {similarity:.4f} - {text}")
        print()
        
        # 测试批量嵌入
        print("3. 批量嵌入测试:")
        batch_embeddings = embedding.get_embeddings(texts[:2])
        print(f"批量处理了 {len(batch_embeddings)} 个文本")
        print(f"每个嵌入向量维度: {len(batch_embeddings[0])}")
        
    except Exception as e:
        print(f"测试过程中发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

